addpath ../
addpath ../export_fig/


% ****************************************************************
% test imputation methods
% 
ntis = 53;
ngene = 100;
nind = 450;
nfactor = 15;
factor_per_gene = 7;
pmiss = .8;

U_true = randn(ntis,nfactor,'single');
Y_full = zeros(ntis,nind,ngene,'single');
Y_train = zeros(ntis,nind,ngene,'single');

factor_map = zeros(ngene,nfactor,'single');

for g = 1:ngene,

    V = randn(factor_per_gene,nind,'single');
    select_factors = randsample(nfactor,factor_per_gene);
    factor_map(g,select_factors) = 1;
    yy = U_true(:,select_factors)*V + 0.1*randn(ntis,nind);
    Y_full(:,:,g) = yy;

    yy(rand(ntis,nind,'single') <= pmiss) = nan;
    Y_train(:,:,g) = yy;
    clear yy;
end


% ****************************************************************
J = 30;
burnin = 400;
inittime = 100;
nepoch = 500;
batchsize = 200;

% first impute the missing values
outdir = 'sim11_result';
! rm -rf sim11_result
run_modQTL_factorization_gibbs(Y_train,J,burnin,inittime,nepoch,batchsize,outdir);


% check imputation accuracy
Y_impute = zeros(ntis,nind,ngene);

load('sim11_result/parameters.mat');
files = dir('sim11_result/V_minibatch*.mat');

for ff = 1:numel(files),
    load([outdir,'/',files(ff).name]);
    for g = 1:numel(genes_minibatch),
        Y_impute(:,:,genes_minibatch(g)) = U * V_minibatch(:,:,g);
    end
    clear *_minibatch;
end

yytrue = Y_full(isnan(Y_train));
yyhat = Y_impute(isnan(Y_train));

[imputation_r,p] = corr( yytrue, yyhat );

close all;
h1 = figure(1);
set(h1,'position',[.5 .5 400 300]);
idx = randsample(numel(yytrue),10000);
hold on;
dscatter( yytrue(idx), yyhat(idx) );
plot( [-3,3], imputation_r*[-3,3], 'r--', 'linewidth', 2 );
plot( [-3,3], [-3,3], 'k--', 'linewidth', 2 );
hold off;
xlabel('true'); ylabel('imputed');
title(sprintf('r = %.2e p = %.2e',imputation_r,p));
axis('tight');

export_fig('sim11_fig01','-pdf','-a1','-m2','-transparent');
